import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.stattools import adfuller
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.seasonal import seasonal_decompose
from sklearn.metrics import mean_squared_error, mean_absolute_error
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体支持
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
matplotlib.rcParams['axes.unicode_minus'] = False

def load_and_preprocess_data():
    """
    加载并预处理数据
    """
    # 读取数据
    base = pd.read_csv('customer_base.csv')
    behavior = pd.read_csv('customer_behavior_assets.csv')
    
    # 转换统计月份为日期格式
    behavior['stat_date'] = pd.to_datetime(behavior['stat_month'], format='%Y-%m')
    
    return base, behavior

def prepare_time_series_data(behavior_data, customer_id=None):
    """
    准备时间序列数据
    如果未指定客户ID，则使用所有客户总体AUM数据
    """
    if customer_id:
        # 分析特定客户的时间序列数据
        customer_data = behavior_data[behavior_data['customer_id'] == customer_id].copy()
        if len(customer_data) < 3:
            print(f"客户 {customer_id} 数据不足，无法进行时间序列分析")
            return None
        customer_data = customer_data.sort_values('stat_date')
        ts_data = customer_data.set_index('stat_date')['total_assets']
    else:
        # 分析所有客户的总体AUM数据
        monthly_aum = behavior_data.groupby('stat_date')['total_assets'].sum()
        ts_data = monthly_aum.sort_index()
    
    return ts_data

def check_stationarity(timeseries, title):
    """
    检查时间序列的平稳性
    """
    print(f'时间序列 {title} 平稳性检验结果:')
    result = adfuller(timeseries.dropna())
    print(f'ADF统计量: {result[0]:.6f}')
    print(f'p值: {result[1]:.6f}')
    
    if result[1] <= 0.05:
        print("序列是平稳的 (拒绝原假设)")
        is_stationary = True
    else:
        print("序列是非平稳的 (接受原假设)")
        is_stationary = False
    
    print('')
    return is_stationary

def difference_series(timeseries, order=1):
    """
    对时间序列进行差分
    """
    diff_series = timeseries.diff(order).dropna()
    return diff_series

def plot_time_series_analysis(ts_data, title):
    """
    绘制时间序列分析图表
    """
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    fig.suptitle(f'{title} 时间序列分析', fontsize=16)
    
    # 原始时间序列
    axes[0, 0].plot(ts_data.index, ts_data.values)
    axes[0, 0].set_title('原始时间序列')
    axes[0, 0].set_xlabel('日期')
    axes[0, 0].set_ylabel('总资产 (AUM)')
    axes[0, 0].tick_params(axis='x', rotation=45)
    
    # 一阶差分
    diff_1 = difference_series(ts_data, 1)
    axes[0, 1].plot(diff_1.index, diff_1.values)
    axes[0, 1].set_title('一阶差分')
    axes[0, 1].set_xlabel('日期')
    axes[0, 1].set_ylabel('差分值')
    axes[0, 1].tick_params(axis='x', rotation=45)
    
    # ACF图
    # 限制lags不超过数据点数-1且不超过数据点数的一半
    max_lags = min(10, len(ts_data) // 2 - 1, len(ts_data) - 1)
    plot_acf(ts_data.dropna(), ax=axes[1, 0], lags=max_lags)
    axes[1, 0].set_title('自相关函数 (ACF)')
    
    # PACF图
    plot_pacf(ts_data.dropna(), ax=axes[1, 1], lags=max_lags)
    axes[1, 1].set_title('偏自相关函数 (PACF)')
    
    plt.tight_layout()
    plt.savefig(f'image_show/{title.replace(" ", "_")}_ts_analysis.png', dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"时间序列分析图表已保存到 {title.replace(' ', '_')}_ts_analysis.png")

def plot_detailed_ts_analysis(ts_data, title):
    """
    绘制详细的时序分析图表，包括趋势、增长率、移动平均
    """
    # 创建图形和子图
    fig = plt.figure(figsize=(20, 10))
    fig.suptitle(f'{title} 详细时序分析', fontsize=20)
    
    # 1. 银行客户总资产月度趋势
    ax1 = plt.subplot(2, 2, 1)
    ax1.plot(ts_data.index, ts_data.values, marker='o', linewidth=2, markersize=6)
    ax1.set_title('银行客户总资产月度趋势', fontsize=14)
    ax1.set_xlabel('日期')
    ax1.set_ylabel('总资产 (AUM)')
    ax1.grid(True, alpha=0.3)
    plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45)
    
    # 2. 总资产月度增长率
    ax2 = plt.subplot(2, 2, 2)
    growth_rate = ts_data.pct_change() * 100
    ax2.plot(growth_rate.index, growth_rate.values, color='orange', marker='s', linewidth=2)
    ax2.set_title('总资产月度增长率 (%)', fontsize=14)
    ax2.set_xlabel('日期')
    ax2.set_ylabel('增长率 (%)')
    ax2.set_ylim(2.075, max(growth_rate.max(), 2.5) + 0.025)  # 设置纵坐标范围
    ax2.set_yticks(np.arange(2.075, max(growth_rate.max(), 2.5) + 0.025, 0.025))  # 设置纵坐标步长
    ax2.grid(True, alpha=0.3)
    plt.setp(ax2.xaxis.get_majorticklabels(), rotation=45)
    
    # 3. 总资产趋势与移动平均
    ax3 = plt.subplot(2, 2, 3)
    ax3.plot(ts_data.index, ts_data.values, label='实际总资产', marker='o')
    # 计算3个月和6个月移动平均
    ma_3 = ts_data.rolling(window=3).mean()
    ma_6 = ts_data.rolling(window=6).mean()
    ax3.plot(ma_3.index, ma_3.values, label='3个月移动平均', linestyle='--', linewidth=2)
    ax3.plot(ma_6.index, ma_6.values, label='6个月移动平均', linestyle='-.', linewidth=2)
    ax3.set_title('总资产趋势与移动平均', fontsize=14)
    ax3.set_xlabel('日期')
    ax3.set_ylabel('总资产 (AUM)')
    ax3.legend()
    ax3.grid(True, alpha=0.3)
    plt.setp(ax3.xaxis.get_majorticklabels(), rotation=45)
    
    # 4. 总资产波动性
    ax4 = plt.subplot(2, 2, 4)
    rolling_std = ts_data.rolling(window=3).std()
    ax4.plot(rolling_std.index, rolling_std.values, color='red', marker='d')
    ax4.set_title('总资产波动性 (3个月滚动标准差)', fontsize=14)
    ax4.set_xlabel('日期')
    ax4.set_ylabel('标准差')
    ax4.grid(True, alpha=0.3)
    plt.setp(ax4.xaxis.get_majorticklabels(), rotation=45)
    
    plt.tight_layout()
    plt.savefig(f'image_show/{title.replace(" ", "_")}_detailed_analysis.png', dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"详细时序分析图表已保存到 {title.replace(' ', '_')}_detailed_analysis.png")

def fit_arima_model(ts_data, order=(1, 1, 1)):
    """
    拟合ARIMA模型
    """
    try:
        model = ARIMA(ts_data, order=order)
        fitted_model = model.fit()
        return fitted_model
    except Exception as e:
        print(f"模型拟合失败: {e}")
        return None

def forecast_aum(ts_data, fitted_model, steps=4):
    """
    预测未来AUM
    """
    # 获取历史数据用于对比
    history = ts_data.copy()
    
    # 进行预测
    forecast_result = fitted_model.forecast(steps=steps)
    forecast_ci = fitted_model.get_forecast(steps=steps).conf_int()
    
    return forecast_result, forecast_ci

def plot_forecast(ts_data, forecast, forecast_ci, title):
    """
    绘制预测结果
    """
    plt.figure(figsize=(12, 6))
    
    # 绘制历史数据
    plt.plot(ts_data.index, ts_data.values, label='历史数据', color='blue')
    
    # 创建未来日期索引
    last_date = ts_data.index[-1]
    if hasattr(last_date, 'to_period'):
        future_dates = pd.date_range(start=last_date + pd.DateOffset(months=1), periods=len(forecast), freq='MS')
    else:
        future_dates = pd.date_range(start=last_date + pd.DateOffset(months=1), periods=len(forecast), freq='MS')
    
    # 绘制预测数据
    plt.plot(future_dates, forecast, label='预测数据', color='red', linestyle='--')
    
    # 绘制置信区间
    plt.fill_between(future_dates, 
                     forecast_ci.iloc[:, 0], 
                     forecast_ci.iloc[:, 1], 
                     color='pink', alpha=0.3, label='95%置信区间')
    
    plt.title(f'{title} AUM预测')
    plt.xlabel('日期')
    plt.ylabel('总资产 (AUM)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xticks(rotation=45)
    
    plt.tight_layout()
    plt.savefig(f'image_show/{title.replace(" ", "_")}_forecast.png', dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"AUM预测图表已保存到 {title.replace(' ', '_')}_forecast.png")

def evaluate_model(ts_data, fitted_model):
    """
    评估模型性能
    """
    # 获取样本内预测
    in_sample_pred = fitted_model.fittedvalues
    
    # 计算评估指标
    mse = mean_squared_error(ts_data.iloc[1:], in_sample_pred.iloc[1:])
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(ts_data.iloc[1:], in_sample_pred.iloc[1:])
    
    print("模型评估指标:")
    print(f"均方误差 (MSE): {mse:.2f}")
    print(f"均方根误差 (RMSE): {rmse:.2f}")
    print(f"平均绝对误差 (MAE): {mae:.2f}")
    print(f"AIC: {fitted_model.aic:.2f}")
    print(f"BIC: {fitted_model.bic:.2f}")

def analyze_individual_customer(base_data, behavior_data, customer_id):
    """
    分析单个客户的时间序列
    """
    print(f"\n=== 客户 {customer_id} AUM时间序列分析 ===")
    
    # 准备时间序列数据
    ts_data = prepare_time_series_data(behavior_data, customer_id)
    if ts_data is None:
        return
    
    print(f"数据点数量: {len(ts_data)}")
    print(f"数据时间范围: {ts_data.index.min()} 至 {ts_data.index.max()}")
    
    # 检查平稳性
    is_stationary = check_stationarity(ts_data, f"客户{customer_id}")
    
    # 绘制时间序列分析图
    plot_time_series_analysis(ts_data, f"客户{customer_id}")
    
    # 拟合ARIMA模型
    print("拟合ARIMA模型...")
    model = fit_arima_model(ts_data, order=(1, 1, 1))
    if model is None:
        print("模型拟合失败")
        return
    
    print("模型摘要:")
    print(model.summary())
    
    # 评估模型
    evaluate_model(ts_data, model)
    
    # 进行预测
    print("\n进行未来4个月AUM预测...")
    forecast, forecast_ci = forecast_aum(ts_data, model, steps=4)
    
    print("预测结果:")
    future_dates = pd.date_range(start=ts_data.index[-1] + pd.DateOffset(months=1), periods=4, freq='MS')
    for i, (date, value) in enumerate(zip(future_dates, forecast)):
        print(f"{date.strftime('%Y-%m')}: {value:,.2f}")
    
    # 绘制预测图
    plot_forecast(ts_data, forecast, forecast_ci, f"客户{customer_id}")

def analyze_overall_aum(behavior_data):
    """
    分析总体AUM时间序列
    """
    print("\n=== 总体AUM时间序列分析 ===")
    
    # 准备时间序列数据
    ts_data = prepare_time_series_data(behavior_data)
    if ts_data is None:
        return
    
    print(f"数据点数量: {len(ts_data)}")
    print(f"数据时间范围: {ts_data.index.min()} 至 {ts_data.index.max()}")
    
    # 检查平稳性
    is_stationary = check_stationarity(ts_data, "总体AUM")
    
    # 绘制时间序列分析图
    plot_time_series_analysis(ts_data, "总体AUM")
    
    # 绘制详细时序分析图
    print("绘制详细时序分析图...")
    plot_detailed_ts_analysis(ts_data, "总体AUM")
    
    # 拟合ARIMA模型
    print("拟合ARIMA模型...")
    model = fit_arima_model(ts_data, order=(1, 1, 1))
    if model is None:
        print("模型拟合失败")
        return
    
    print("模型摘要:")
    print(model.summary())
    
    # 评估模型
    evaluate_model(ts_data, model)
    
    # 进行预测
    print("\n进行未来4个月总体AUM预测...")
    forecast, forecast_ci = forecast_aum(ts_data, model, steps=4)
    
    print("预测结果:")
    future_dates = pd.date_range(start=ts_data.index[-1] + pd.DateOffset(months=1), periods=4, freq='MS')
    for i, (date, value) in enumerate(zip(future_dates, forecast)):
        print(f"{date.strftime('%Y-%m')}: {value:,.2f}")
    
    # 绘制预测图
    plot_forecast(ts_data, forecast, forecast_ci, "总体AUM")

def main():
    """
    主函数
    """
    print("开始AUM时间序列分析...")
    
    # 加载数据
    print("1. 加载数据...")
    base_data, behavior_data = load_and_preprocess_data()
    print(f"基础数据加载完成，共 {len(base_data)} 条客户记录")
    print(f"行为数据加载完成，共 {len(behavior_data)} 条记录")
    
    # 分析总体AUM趋势
    analyze_overall_aum(behavior_data)
    
    # 选择几个客户进行个体分析
    sample_customers = behavior_data['customer_id'].unique()[:3]
    for customer_id in sample_customers:
        analyze_individual_customer(base_data, behavior_data, customer_id)
    
    print("\nAUM时间序列分析完成！")

if __name__ == '__main__':
    main()